#pragma once

enum Colour
{
	RED,
	BLACK
};


template<class K, class V>
struct RBTreeNode
{
	RBTreeNode<K, V>* _left;
	RBTreeNode<K, V>* _right;
	RBTreeNode<K, V>* _parent;
	pair<K, V> _kv;
	Colour _col;

	RBTreeNode(const pair<K, V>& kv)
		:_left(nullptr)
		, _right(nullptr)
		, _parent(nullptr)
		, _kv(kv)
		, _col(RED)
	{}
};


template<class K, class V>
class RBTree
{
	typedef RBTreeNode<K, V> Node;
public:
	bool Insert(const pair<K, V>& kv);
	void InOrder();
	bool IsBalance();
	int Height();
	int Size();
	Node* Find(const K& key);
	

private:
	void RotateL(Node* parent);
	void RotateR(Node* parent);
	void _InOrder(Node* root);
	bool _Check(Node* root, int blacknum, const int refVal);
	int _Height(Node* root);
	int _Size(Node* root);

	Node* _root = nullptr;
};




template<class K, class V>
bool RBTree<K,V>::Insert(const pair<K, V>& kv)
{
	if (_root == nullptr)
	{
		_root = new Node(kv);
		_root->_col = BLACK;
		return true;
	}

	Node* parent = nullptr;
	Node* cur = _root;

	while (cur)
	{
		if (cur->_kv.first < kv.first)
		{
			parent = cur;
			cur = cur->_right;
		}
		else if (cur->_kv.first > kv.first)
		{
			parent = cur;
			cur = cur->_left;
		}
		else
		{
			return false;
		}
	}

	// 新增节点给红色
	cur = new Node(kv);
	cur->_col = RED;
	if (parent->_kv.first < kv.first)
	{
		parent->_right = cur;
		cur->_parent = parent;
	}
	else
	{
		parent->_left = cur;
		cur->_parent = parent;
	}

	while (parent && parent->_col == RED)
	{
		Node* grandfather = parent->_parent;
		if (parent == grandfather->_left)
		{
			//     g
			//   p   u
			// c
			Node* uncle = grandfather->_right;
			if (uncle && uncle->_col == RED)
			{
				// 变色
				parent->_col = uncle->_col = BLACK;
				grandfather->_col = RED;

				// 继续往上更新处理
				cur = grandfather;
				parent = cur->_parent;
			}
			else
			{
				if (cur == parent->_left)
				{
					// 单旋
					//     g
					//   p
					// c
					RotateR(grandfather);
					parent->_col = BLACK;
					grandfather->_col = RED;
				}
				else
				{
					// 双旋
					//     g
					//   p
					//     c
					RotateL(parent);
					RotateR(grandfather);
					cur->_col = BLACK;
					grandfather->_col = RED;
				}

				break;
			}
		}
		else  // parent == grandfather->_right
		{
			//     g
			//   u   p 
			//          c
			//
			Node* uncle = grandfather->_left;
			if (uncle && uncle->_col == RED)
			{
				// 变色
				parent->_col = uncle->_col = BLACK;
				grandfather->_col = RED;

				// 继续往上处理
				cur = grandfather;
				parent = cur->_parent;
			}
			else
			{
				if (cur == parent->_right)
				{
					RotateL(grandfather);
					parent->_col = BLACK;
					grandfather->_col = RED;
				}
				else
				{
					//     g
					//   u   p 
					//     c
					//
					RotateR(parent);
					RotateL(grandfather);
					cur->_col = BLACK;
					grandfather->_col = RED;
				}

				break;
			}
		}
	}

	_root->_col = BLACK;

	return true;
}

template<class K, class V>
void RBTree<K,V>::RotateL(Node* parent)
{
	Node* subR = parent->_right;
	Node* subRL = subR->_left;

	parent->_right = subRL;
	subR->_left = parent;

	Node* parentParent = parent->_parent;

	parent->_parent = subR;
	if (subRL)
		subRL->_parent = parent;

	if (_root == parent)
	{
		_root = subR;
		subR->_parent = nullptr;
	}
	else
	{
		if (parentParent->_left == parent)
		{
			parentParent->_left = subR;
		}
		else
		{
			parentParent->_right = subR;
		}

		subR->_parent = parentParent;
	}
}

template<class K, class V>
void RBTree<K,V>::RotateR(Node* parent)
{
	Node* subL = parent->_left;
	Node* subLR = subL->_right;

	parent->_left = subLR;
	if (subLR)
		subLR->_parent = parent;

	Node* parentParent = parent->_parent;

	subL->_right = parent;
	parent->_parent = subL;

	if (_root == parent)
	{
		_root = subL;
		subL->_parent = nullptr;
	}
	else
	{
		if (parentParent->_left == parent)
		{
			parentParent->_left = subL;
		}
		else
		{
			parentParent->_right = subL;
		}

		subL->_parent = parentParent;
	}
}

template<class K, class V>
void RBTree<K,V>::InOrder()
{
	_InOrder(_root);
	cout << endl;
}

template<class K, class V>
void RBTree<K, V>::_InOrder(Node* root)
{
	if (root == nullptr)
		return;

	_InOrder(root->_left);
	cout << root->_kv.first << " ";
	_InOrder(root->_right);
}

// 根节点->当前节点这条路径的黑色节点的数量
template<class K, class V>
bool RBTree<K, V>::_Check(Node* root, int blacknum, const int refVal)
{
	if (root == nullptr)
	{
		//cout << balcknum << endl;
		if (blacknum != refVal)
		{
			cout << "存在黑色节点数量不相等的路径" << endl;
			return false;
		}

		return true;
	}

	if (root->_col == RED && root->_parent->_col == RED)
	{
		cout << "有连续的红色节点" << endl;

		return false;
	}

	if (root->_col == BLACK)
	{
		++blacknum;
	}

	return _Check(root->_left, blacknum, refVal)
		&& _Check(root->_right, blacknum, refVal);
}

template<class K, class V>
bool RBTree<K, V>::IsBalance()
{
	if (_root == nullptr)
		return true;

	if (_root->_col == RED)
		return false;

	//参考值
	int refVal = 0;
	Node* cur = _root;
	while (cur)
	{
		if (cur->_col == BLACK)
		{
			++refVal;
		}

		cur = cur->_left;
	}

	int blacknum = 0;
	return _Check(_root, blacknum, refVal);
}

template<class K, class V>
int RBTree<K,V>::Height()
{
	return _Height(_root);
}

template<class K, class V>
int RBTree<K, V>::_Height(Node* root)
{
	if (root == nullptr)
		return 0;

	int leftHeight = _Height(root->_left);
	int rightHeight = _Height(root->_right);

	return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
}

template<class K, class V>
int RBTree<K, V>::Size()
{
	return _Size(_root);
}

template<class K, class V>
int RBTree<K, V>::_Size(Node* root)
{
	return root == nullptr ? 0 :
		_Size(root->_left) +
		_Size(root->_right) + 1;
}

template<class K, class V>
RBTreeNode<K, V>* RBTree<K,V>::Find(const K& key)
{
	Node* cur = _root;
	while (cur)
	{
		if (cur->_kv.first < key)
		{
			cur = cur->_right;
		}
		else if (cur->_kv.first > key)
		{
			cur = cur->_left;
		}
		else
		{
			return cur;
		}
	}

	return nullptr;
}